#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GitCode&文心大模型&智源研究院AI应用开发大赛
赛道二（空手上场）- 基于ERNIE-4.5-0.3B的智能童话书应用

核心功能：
1. 智能童话故事生成（基于ERNIE-4.5-0.3B）
2. 互动式剧情分支选择
3. 教育知识融合（智源数据增强）
4. 个性化学习伴侣
5. 沉浸式阅读体验
"""

from flask import Flask, render_template, request, jsonify, session
import json
import random
import os
import logging
from datetime import datetime
import uuid

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.secret_key = 'fairy_tale_secret_key_2024'

class ERNIEFairyTaleModel:
    """ERNIE-4.5-0.3B 童话故事生成模型"""
    
    def __init__(self):
        self.model_path = "./ERNIE-4.5-0.3B-Paddle"
        self.is_loaded = False
        self.story_templates = self._load_story_templates()
        self.educational_content = self._load_educational_content()
        
    def _load_story_templates(self):
        """加载童话故事模板"""
        return {
            "adventure": {
                "title": "奇幻冒险",
                "opening": "在一个充满魔法的森林里，住着一位勇敢的小{character}...",
                "themes": ["勇气", "友谊", "探索", "成长"],
                "characters": ["小兔子", "小熊", "小鸟", "小松鼠", "小狐狸"]
            },
            "friendship": {
                "title": "友谊故事", 
                "opening": "在美丽的花园里，{character1}和{character2}是最好的朋友...",
                "themes": ["友谊", "分享", "互助", "理解"],
                "characters": ["小猫", "小狗", "小羊", "小猪", "小鸭"]
            },
            "learning": {
                "title": "学习成长",
                "opening": "聪明的小{character}今天要学习一个新本领...",
                "themes": ["学习", "坚持", "智慧", "创新"],
                "characters": ["小象", "小猴子", "小海豚", "小企鹅", "小老鼠"]
            },
            "nature": {
                "title": "自然探索",
                "opening": "在神奇的大自然中，小{character}发现了一个秘密...",
                "themes": ["环保", "自然", "科学", "观察"],
                "characters": ["小蜜蜂", "小蝴蝶", "小青蛙", "小鱼", "小鸟"]
            }
        }
    
    def _load_educational_content(self):
        """加载教育内容（模拟智源数据集）"""
        return {
            "science": {
                "physics": ["光的传播", "声音的产生", "磁铁的奥秘", "水的三态变化"],
                "biology": ["植物的生长", "动物的习性", "生态系统", "食物链"],
                "chemistry": ["颜色的变化", "物质的状态", "简单反应", "日常化学"]
            },
            "math": {
                "numbers": ["数字游戏", "简单计算", "图形认知", "逻辑推理"],
                "geometry": ["形状识别", "空间概念", "对称美学", "测量概念"]
            },
            "language": {
                "vocabulary": ["词汇扩展", "成语故事", "诗词韵律", "表达技巧"],
                "reading": ["阅读理解", "故事结构", "情感表达", "想象力"]
            },
            "social": {
                "values": ["诚实守信", "尊重他人", "团队合作", "责任担当"],
                "culture": ["传统文化", "节日习俗", "地理知识", "历史故事"]
            }
        }
    
    def load_model(self):
        """加载ERNIE模型"""
        try:
            if os.path.exists(self.model_path):
                logger.info(f"正在加载ERNIE-4.5-0.3B模型: {self.model_path}")
                # 这里应该是实际的模型加载代码
                # import paddle
                # self.model = paddle.jit.load(self.model_path)
                self.is_loaded = True
                logger.info("ERNIE模型加载成功")
                return True
            else:
                logger.warning("ERNIE模型路径不存在，使用模拟模式")
                self.is_loaded = False
                return False
        except Exception as e:
            logger.error(f"模型加载失败: {e}")
            self.is_loaded = False
            return False
    
    def generate_story(self, theme, character_name, user_choices=None, educational_focus=None):
        """生成童话故事"""
        try:
            if theme not in self.story_templates:
                theme = "adventure"
            
            template = self.story_templates[theme]
            characters = template["characters"]
            main_character = character_name or random.choice(characters)
            
            # 基础故事开头
            if theme == "friendship":
                character2 = random.choice([c for c in characters if c != main_character])
                opening = template["opening"].format(character1=main_character, character2=character2)
            else:
                opening = template["opening"].format(character=main_character)
            
            # 生成故事内容
            story_parts = [opening]
            
            # 添加教育元素
            if educational_focus:
                edu_content = self._get_educational_content(educational_focus)
                story_parts.append(f"\n在这个过程中，{main_character}学到了{edu_content}的知识。")
            
            # 根据用户选择发展剧情
            if user_choices:
                for choice in user_choices:
                    story_parts.append(self._develop_plot(choice, main_character, theme))
            
            # 生成结尾
            ending = self._generate_ending(theme, main_character)
            story_parts.append(ending)
            
            # 组合完整故事
            full_story = "".join(story_parts)
            
            # 生成选择分支
            choices = self._generate_choices(theme, main_character)
            
            return {
                "story": full_story,
                "choices": choices,
                "theme": template["title"],
                "character": main_character,
                "educational_value": self._extract_educational_value(theme, educational_focus)
            }
            
        except Exception as e:
            logger.error(f"故事生成失败: {e}")
            return self._get_fallback_story()
    
    def _get_educational_content(self, focus):
        """获取教育内容"""
        if focus in self.educational_content:
            category = random.choice(list(self.educational_content[focus].keys()))
            content = random.choice(self.educational_content[focus][category])
            return content
        return "友谊与合作"
    
    def _develop_plot(self, choice, character, theme):
        """根据选择发展剧情"""
        plot_developments = {
            "brave": f"\n{character}鼓起勇气，决定面对挑战。在勇敢的行动中，{character}发现自己比想象中更强大。",
            "help": f"\n{character}选择帮助遇到困难的朋友。通过互相帮助，他们一起解决了问题。",
            "explore": f"\n{character}决定探索未知的地方。在探索过程中，{character}发现了许多有趣的秘密。",
            "learn": f"\n{character}选择学习新的知识。通过努力学习，{character}掌握了新的技能。",
            "share": f"\n{character}决定与朋友分享自己的发现。分享让快乐变得更加美好。",
            "think": f"\n{character}停下来仔细思考。智慧的思考帮助{character}找到了最好的解决方案。"
        }
        return plot_developments.get(choice, f"\n{character}做出了明智的选择，故事继续发展...")
    
    def _generate_ending(self, theme, character):
        """生成故事结尾"""
        endings = {
            "adventure": f"\n\n经过这次奇妙的冒险，{character}不仅变得更加勇敢，还收获了珍贵的友谊和宝贵的经验。从此，{character}成为了森林里最受欢迎的小英雄。",
            "friendship": f"\n\n通过这次经历，{character}和朋友们的友谊变得更加深厚。他们明白了真正的友谊需要理解、包容和互相帮助。",
            "learning": f"\n\n{character}通过不断的学习和努力，终于掌握了新的本领。这让{character}明白了学习的重要性和坚持的力量。",
            "nature": f"\n\n{character}在大自然中学到了很多知识，也懂得了保护环境的重要性。从此，{character}成为了大自然的小守护者。"
        }
        return endings.get(theme, f"\n\n{character}的故事告诉我们，每个人都有自己的闪光点，只要勇敢地去追求梦想，就一定能够成功。")
    
    def _generate_choices(self, theme, character):
        """生成选择分支"""
        base_choices = [
            {"id": "brave", "text": f"让{character}勇敢面对挑战", "theme": "勇气"},
            {"id": "help", "text": f"让{character}帮助其他小动物", "theme": "友善"},
            {"id": "explore", "text": f"让{character}探索新的地方", "theme": "探索"},
            {"id": "learn", "text": f"让{character}学习新知识", "theme": "学习"},
            {"id": "share", "text": f"让{character}分享快乐", "theme": "分享"},
            {"id": "think", "text": f"让{character}仔细思考", "theme": "智慧"}
        ]
        
        # 根据主题筛选相关选择
        theme_relevant = {
            "adventure": ["brave", "explore", "help"],
            "friendship": ["help", "share", "think"],
            "learning": ["learn", "think", "explore"],
            "nature": ["explore", "learn", "help"]
        }
        
        relevant_ids = theme_relevant.get(theme, ["brave", "help", "learn"])
        return [choice for choice in base_choices if choice["id"] in relevant_ids]
    
    def _extract_educational_value(self, theme, focus):
        """提取教育价值"""
        values = {
            "adventure": ["勇气与坚持", "面对困难的态度", "团队合作的重要性"],
            "friendship": ["友谊的珍贵", "理解与包容", "互相帮助"],
            "learning": ["学习的乐趣", "知识的力量", "持续进步"],
            "nature": ["环境保护", "科学探索", "自然规律"]
        }
        
        base_values = values.get(theme, ["积极向上", "品德培养"])
        
        if focus:
            focus_values = {
                "science": ["科学思维", "观察能力", "实验精神"],
                "math": ["逻辑思维", "数学概念", "问题解决"],
                "language": ["语言表达", "阅读理解", "创意思维"],
                "social": ["社会责任", "文化传承", "价值观念"]
            }
            base_values.extend(focus_values.get(focus, []))
        
        return base_values
    
    def _get_fallback_story(self):
        """备用故事"""
        return {
            "story": "在一个美丽的森林里，住着一只聪明的小兔子。小兔子每天都会学习新的知识，帮助森林里的朋友们。通过不断的努力和学习，小兔子成为了大家最喜欢的小老师。",
            "choices": [
                {"id": "brave", "text": "让小兔子勇敢探险", "theme": "勇气"},
                {"id": "help", "text": "让小兔子帮助朋友", "theme": "友善"},
                {"id": "learn", "text": "让小兔子学习新知识", "theme": "学习"}
            ],
            "theme": "成长故事",
            "character": "小兔子",
            "educational_value": ["友谊", "学习", "成长"]
        }

class ZhiyuanDataEnhancer:
    """智源数据增强器"""
    
    def __init__(self):
        self.infinity_instruct_data = self._load_infinity_instruct_sample()
    
    def _load_infinity_instruct_sample(self):
        """加载Infinity-Instruct数据样本"""
        return {
            "educational_qa": [
                {"question": "为什么天空是蓝色的？", "answer": "天空呈现蓝色是因为阳光中的蓝色光波长较短，更容易被大气中的小颗粒散射。"},
                {"question": "植物为什么需要阳光？", "answer": "植物需要阳光进行光合作用，将二氧化碳和水转化为养分和氧气。"},
                {"question": "为什么会有四季变化？", "answer": "四季变化是因为地球绕太阳公转时，地轴倾斜导致不同地区接受阳光的角度和时间不同。"}
            ],
            "story_elements": [
                {"element": "勇气", "description": "面对困难时不退缩，勇敢地去尝试和挑战。"},
                {"element": "友谊", "description": "真诚地对待朋友，互相帮助，共同成长。"},
                {"element": "智慧", "description": "善于思考，用知识和经验解决问题。"}
            ],
            "learning_tips": [
                {"skill": "观察力", "tip": "仔细观察周围的事物，发现其中的规律和变化。"},
                {"skill": "想象力", "tip": "大胆想象，创造属于自己的奇妙世界。"},
                {"skill": "表达力", "tip": "用清晰的语言表达自己的想法和感受。"}
            ]
        }
    
    def enhance_story_with_knowledge(self, story_data, knowledge_type="educational_qa"):
        """用智源知识增强故事"""
        if knowledge_type in self.infinity_instruct_data:
            knowledge_item = random.choice(self.infinity_instruct_data[knowledge_type])
            story_data["knowledge_enhancement"] = knowledge_item
        return story_data
    
    def get_learning_suggestions(self, theme):
        """获取学习建议"""
        suggestions = {
            "adventure": ["培养勇气和探索精神", "学习解决问题的方法", "了解团队合作的重要性"],
            "friendship": ["学会理解和包容他人", "培养同理心", "掌握沟通技巧"],
            "learning": ["培养好奇心和求知欲", "学会坚持和努力", "掌握学习方法"],
            "nature": ["了解自然科学知识", "培养环保意识", "学会观察和思考"]
        }
        return suggestions.get(theme, ["培养积极的人生态度", "学会思考和表达", "建立正确的价值观"])

# 初始化模型和数据增强器
ernie_model = ERNIEFairyTaleModel()
zhiyuan_enhancer = ZhiyuanDataEnhancer()

# 启动时加载模型
ernie_model.load_model()

@app.route('/')
def index():
    """主页"""
    return render_template('fairy_tale_index.html')

@app.route('/api/generate_story', methods=['POST'])
def generate_story():
    """生成童话故事API"""
    try:
        data = request.get_json()
        theme = data.get('theme', 'adventure')
        character_name = data.get('character_name', '')
        user_choices = data.get('choices', [])
        educational_focus = data.get('educational_focus', None)
        
        # 生成故事
        story_data = ernie_model.generate_story(
            theme=theme,
            character_name=character_name,
            user_choices=user_choices,
            educational_focus=educational_focus
        )
        
        # 用智源数据增强
        enhanced_story = zhiyuan_enhancer.enhance_story_with_knowledge(story_data)
        
        # 添加学习建议
        enhanced_story["learning_suggestions"] = zhiyuan_enhancer.get_learning_suggestions(theme)
        
        # 保存到会话
        if 'story_history' not in session:
            session['story_history'] = []
        
        session['story_history'].append({
            'id': str(uuid.uuid4()),
            'timestamp': datetime.now().isoformat(),
            'story_data': enhanced_story
        })
        
        return jsonify({
            'success': True,
            'data': enhanced_story
        })
        
    except Exception as e:
        logger.error(f"故事生成API错误: {e}")
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/api/continue_story', methods=['POST'])
def continue_story():
    """继续故事API"""
    try:
        data = request.get_json()
        choice_id = data.get('choice_id')
        current_story = data.get('current_story', '')
        theme = data.get('theme', 'adventure')
        character = data.get('character', '小主角')
        
        # 根据选择继续故事
        continuation = ernie_model._develop_plot(choice_id, character, theme)
        
        # 生成新的选择
        new_choices = ernie_model._generate_choices(theme, character)
        
        return jsonify({
            'success': True,
            'data': {
                'continuation': continuation,
                'choices': new_choices
            }
        })
        
    except Exception as e:
        logger.error(f"故事继续API错误: {e}")
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/api/get_story_history')
def get_story_history():
    """获取故事历史"""
    try:
        history = session.get('story_history', [])
        return jsonify({
            'success': True,
            'data': history
        })
    except Exception as e:
        logger.error(f"获取历史API错误: {e}")
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/api/system_status')
def system_status():
    """系统状态API"""
    return jsonify({
        'status': 'running',
        'model_loaded': ernie_model.is_loaded,
        'model_path': ernie_model.model_path,
        'zhiyuan_data_loaded': True,
        'timestamp': datetime.now().isoformat(),
        'app_name': '智能童话书竞赛应用',
        'competition': 'GitCode&文心大模型&智源研究院AI应用开发大赛',
        'track': '赛道二（空手上场）'
    })

@app.route('/api/educational_content')
def get_educational_content():
    """获取教育内容API"""
    try:
        return jsonify({
            'success': True,
            'data': ernie_model.educational_content
        })
    except Exception as e:
        logger.error(f"教育内容API错误: {e}")
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

if __name__ == '__main__':
    print("🎭 智能童话书竞赛应用启动中...")
    print("📚 基于ERNIE-4.5-0.3B + 智源研究院数据")
    print("🏆 GitCode&文心大模型&智源研究院AI应用开发大赛")
    print("🌟 赛道二（空手上场）")
    print("=" * 50)
    print("🚀 应用地址: http://localhost:5000")
    print("=" * 50)
    
    app.run(debug=True, host='0.0.0.0', port=5000)